#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 11:00:30
# @LastEditTime: 2025-03-14 11:00:35
# @Description: 告警抑制模块
###

import json
import time
import redis
from typing import Dict, Optional
from datetime import datetime, timedelta

class AlertSuppression:
    """告警抑制管理器"""
    
    def __init__(self, redis_config: Dict):
        """初始化告警抑制器
        
        Args:
            redis_config: Redis配置
        """
        self.redis = redis.Redis(
            host=redis_config.get('host', 'localhost'),
            port=redis_config.get('port', 6379),
            password=redis_config.get('password'),
            db=redis_config.get('db', 0)
        )
        # 告警抑制时间（秒）
        self.suppression_time = {
            'critical': 300,   # 严重告警 5分钟
            'warning': 600,    # 警告 10分钟
            'info': 1800       # 信息 30分钟
        }
        
    def _get_alert_key(self, alert: Dict) -> str:
        """生成告警的唯一键
        
        Args:
            alert: 告警信息
        """
        # 生成告警的唯一标识
        key_parts = [
            alert.get('platform', ''),           # 云平台
            alert.get('resource_type', ''),      # 资源类型
            alert.get('resource_id', ''),        # 资源ID
            alert.get('metric', ''),             # 指标名称
            alert.get('alert_type', 'warning'),  # 告警类型
        ]
        return ':'.join(filter(None, key_parts))
        
    def should_suppress(self, alert: Dict) -> bool:
        """判断是否需要抑制告警
        
        Args:
            alert: 告警信息
        """
        alert_key = self._get_alert_key(alert)
        
        # 获取上次告警时间
        last_alert = self.redis.get(alert_key)
        if not last_alert:
            return False
            
        last_alert_data = json.loads(last_alert)
        last_time = datetime.fromtimestamp(last_alert_data['timestamp'])
        alert_type = alert.get('alert_type', 'warning')
        
        # 计算是否在抑制时间内
        suppression_seconds = self.suppression_time.get(alert_type, 600)
        if datetime.now() - last_time < timedelta(seconds=suppression_seconds):
            # 更新告警次数
            last_alert_data['count'] += 1
            self.redis.set(
                alert_key,
                json.dumps(last_alert_data),
                ex=suppression_seconds
            )
            return True
            
        return False
        
    def record_alert(self, alert: Dict) -> None:
        """记录告警
        
        Args:
            alert: 告警信息
        """
        alert_key = self._get_alert_key(alert)
        alert_type = alert.get('alert_type', 'warning')
        
        # 记录告警信息
        alert_data = {
            'timestamp': time.time(),
            'count': 1,
            'alert': alert
        }
        
        # 设置过期时间
        self.redis.set(
            alert_key,
            json.dumps(alert_data),
            ex=self.suppression_time.get(alert_type, 600)
        )
        
    def get_alert_count(self, alert: Dict) -> int:
        """获取告警次数
        
        Args:
            alert: 告警信息
        """
        alert_key = self._get_alert_key(alert)
        last_alert = self.redis.get(alert_key)
        
        if last_alert:
            last_alert_data = json.loads(last_alert)
            return last_alert_data.get('count', 0)
            
        return 0
